from typing import TYPE_CHECKING
import matplotlib.pyplot as plt

from puzzle.mathF import get_center
from container import Tiles, OPTIONS, v_get_face_color_by_td
from visualizer.board_visualizer import visualize_board
import visualizer.common as v

if TYPE_CHECKING:
    from solver.solver import SolDesc
    from solver.SBoard import SBoard
    from matplotlib.axes import Axes
    from matplotlib.figure import Figure

_anim_fig: "Figure" = None # type: ignore
_anim_ax: "Axes" = None # type: ignore

def anim_prepare(sb: "SBoard"):
    plt.ion()
    
    fig, ax = plt.subplots()
    v.prepare_ax(ax)
    
    global _anim_fig, _anim_ax
    _anim_fig = fig
    _anim_ax = ax
    
    v.plot_board(_anim_ax, sb.board)
    
def anim_update_sd(sb: "SBoard", sd: "SolDesc", it: int):
    _anim_fig.canvas.manager.set_window_title(f"Iteration {it}") # type: ignore
    
    for row in range(len(sd.bd)):
        for col in range(len(sd.bd[row])):
            td = sb.board[row, col]
            match sd.bd[row][col]:
                case -2: # Blank
                    edgecolor = (1, 1, 1)
                case -1: # Avail
                    edgecolor = (0, 0, 0)
                case _:
                    edgecolor = OPTIONS.V_COMP_COLOR[sd.bd[row][col]]
            
            v.draw_polygon(
                _anim_ax,
                td.vertices,
                edgecolor,
            )
    
    plt.draw()
    plt.pause(0.2)

def anim_stop():
    plt.ioff()

def visualize_sd(sb: "SBoard", sd: "SolDesc") -> "Axes":
    board = sb.board
    
    ax: "Axes"
    fig, ax = plt.subplots()
    
    pi = []
    for i in range(len(sd.piece_used)):
        if sd.piece_used[i]: pi.append(str(i))
        
    title = f"{sd.n_pieces_used} used: {" ".join(pi)}"
    fig.canvas.manager.set_window_title(title) # type: ignore
    
    v.prepare_ax(ax)

    for td in board.tiles_flatten:
        if td.tile == Tiles.blank:
            continue
        
        center = get_center(td.vertices)
        ax.text(center.x, center.y, td.tile.label, ha="center", va="center")
    
    for row in range(len(sd.bd)):
        for col in range(len(sd.bd[row])):
            td = board[row, col]
            
            if sd.bd[row][col] >= 0:
                
                v.draw_polygon(
                    ax,
                    td.vertices,
                    facecolor = OPTIONS.V_COMP_COLOR[sd.bd[row][col]],
                    edgecolor = (0, 0, 0),
                    linewidth = 2.0
                )
    
    return ax